package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"
)

// https://kebingzao.com/2020/12/29/go-check-certs/
const defaultConcurrency = 8

type sigAlgSunset struct {
	name      string
	sunsetsAt time.Time
}

const (
	errExpiringShortly = "%s: ** '%s' (S/N %X) 过期 in %d hours! **"
	errExpiringSoon    = "%s: '%s' (S/N %X) 过期 in %d days."
	errSunsetAlg       = "%s: '%s' (S/N %X) 过期 after the sunset date for its signature algorithm '%s'."
)

var sunsetSigAlgs = map[x509.SignatureAlgorithm]sigAlgSunset{
	x509.MD2WithRSA: {
		name:      "MD2 with RSA",
		sunsetsAt: time.Now(),
	},
	x509.MD5WithRSA: {
		name:      "MD5 with RSA",
		sunsetsAt: time.Now(),
	},
	x509.SHA1WithRSA: {
		name:      "SHA1 with RSA",
		sunsetsAt: time.Date(2017, 1, 1, 0, 0, 0, 0, time.UTC),
	},
	x509.DSAWithSHA1: {
		name:      "DSA with SHA1",
		sunsetsAt: time.Date(2017, 1, 1, 0, 0, 0, 0, time.UTC),
	},
	x509.ECDSAWithSHA1: {
		name:      "ECDSA with SHA1",
		sunsetsAt: time.Date(2017, 1, 1, 0, 0, 0, 0, time.UTC),
	},
}

var (
	hostsFile   = flag.String("hosts", "./host.txt", "域名列表")
	warnYears   = flag.Int("years", 0, "年")
	warnMonths  = flag.Int("months", 0, "月")
	warnDays    = flag.Int("days", 0, "日")
	checkSigAlg = flag.Bool("check-sig-alg", true, "Verify that non-root certificates are using a good signature algorithm.")
	concurrency = flag.Int("concurrency", defaultConcurrency, "Maximum number of hosts to check at once.")
)

type certErrors struct {
	commonName string
	errs       []error
}

type hostResult struct {
	host  string
	err   error
	certs []certErrors
}

func fileExists(filename string) bool {
	_, err := os.Stat(filename)
	if err == nil {
		return true
	}
	if os.IsNotExist(err) {
		return false
	}
	// 如果发生其他错误，可以在这里处理
	return false
}

func main() {
	flag.Parse()

	if len(*hostsFile) == 0 {
		flag.Usage()
		return
	}

	if !fileExists(*hostsFile) {
		fmt.Printf("%s 不存在\n", *hostsFile)

		flag.Usage()
		return
	}

	if *warnYears < 0 {
		*warnYears = 0
	}
	if *warnMonths < 0 {
		*warnMonths = 0
	}
	if *warnDays < 0 {
		*warnDays = 0
	}

	if *warnYears == 0 && *warnMonths == 0 && *warnDays == 0 {
		*warnDays = 30
	}
	if *concurrency < 0 {
		*concurrency = defaultConcurrency
	}

	processHosts()
}

func processHosts() {
	done := make(chan struct{})
	defer close(done)
	hosts := queueHosts(done)
	results := make(chan hostResult)

	var wg sync.WaitGroup
	wg.Add(*concurrency)

	for i := 0; i < *concurrency; i++ {
		go func() {
			processQueue(done, hosts, results)
			wg.Done()
		}()
	}
	go func() {
		wg.Wait()
		close(results)
	}()

	for r := range results {
		if r.err != nil {
			// i/o timeout
			// log.Printf("%s: %v\n", r.host, r.err)
			continue
		}

		for _, cert := range r.certs {
			for _, err := range cert.errs {
				log.Printf("%v\n", err)
			}
		}
	}
}

func queueHosts(done <-chan struct{}) <-chan string {
	hosts := make(chan string)
	go func() {
		defer close(hosts)
		fileContents, err := os.ReadFile(*hostsFile)
		if err != nil {
			return
		}
		lines := strings.Split(string(fileContents), "\n")
		for _, line := range lines {
			host := strings.TrimSpace(line)
			if len(host) == 0 || host[0] == '#' {
				continue
			}

			select {
			case hosts <- host:
			case <-done:
				return
			}
		}
	}()
	return hosts
}

func processQueue(done <-chan struct{}, hosts <-chan string, results chan<- hostResult) {
	for host := range hosts {
		select {
		case results <- checkHost(host):
		case <-done:
			return
		}
	}
}

func checkHost(host string) (result hostResult) {
	result = hostResult{
		host:  host,
		certs: []certErrors{},
	}
	// conn, err := tls.Dial("tcp", host, nil)
	timeout := 10 * time.Second // 设置超时时间为10秒
	dialer := &net.Dialer{
		Timeout: timeout,
	}

	conn, err := tls.DialWithDialer(dialer, "tcp", host, nil)
	if err != nil {
		result.err = err
		return
	}
	defer conn.Close()
	timeNow := time.Now()
	checkedCerts := make(map[string]struct{})
	for _, chain := range conn.ConnectionState().VerifiedChains {
		for cerNum, cert := range chain {
			if _, checked := checkedCerts[string(cert.Signature)]; checked {
				continue
			}
			checkedCerts[string(cert.Signature)] = struct{}{}
			cErrs := []error{}

			if timeNow.AddDate(*warnYears, *warnMonths, *warnDays).After(cert.NotAfter) {
				expiresIn := int64(cert.NotAfter.Sub(timeNow).Hours())
				if expiresIn <= 48 {
					cErrs = append(cErrs,
						fmt.Errorf(
							errExpiringShortly,
							host,
							cert.Subject.CommonName,
							cert.SerialNumber,
							expiresIn))
				} else {
					cErrs = append(cErrs,
						fmt.Errorf(
							errExpiringSoon,
							host,
							cert.Subject.CommonName,
							cert.SerialNumber,
							expiresIn/24,
						))
				}

				if alg, exists := sunsetSigAlgs[cert.SignatureAlgorithm]; *checkSigAlg &&
					exists && cerNum != len(chain)-1 {
					if cert.NotAfter.Equal(alg.sunsetsAt) || cert.NotAfter.After(alg.sunsetsAt) {
						cErrs = append(cErrs,
							fmt.Errorf(
								errSunsetAlg,
								host,
								cert.Subject.CommonName,
								cert.SerialNumber,
								alg.name,
							))
					}
				}

				result.certs = append(result.certs, certErrors{
					commonName: cert.Subject.CommonName,
					errs:       cErrs,
				})
			}
		}
	}
	return
}
